import pickle
import cv2
import os
import numpy as np


def write_landmarks_on_img(landmarks, input_img_path, output_img_path):
    input_img = cv2.imread(input_img_path)
    output_img = input_img.copy()
    for i, coordinate in enumerate(landmarks):
        cv2.putText(output_img, str(i), tuple(coordinate.astype(int)), cv2.FONT_HERSHEY_DUPLEX, 0.3, (0, 0, 255), 1)

    cv2.imwrite(output_img_path, output_img)


def find_the_overlapping_idx(landmarks, input_img_path, output_img_path):
    idx1 = 55
    input_img = cv2.imread(input_img_path)
    output_img = input_img.copy()
    distance = np.linalg.norm(landmarks - landmarks[idx1], axis=1)
    # min_idx = np.argmin(distance, axis=0)
    sort_result = np.argsort(distance)
    print("The index is", sort_result[1])
    print("sort_result is", sort_result)
    for i, coordinate in enumerate(landmarks):
        if i == idx1 or i == sort_result[0]:
            # (blue, green, red)
            cv2.putText(output_img, str(i), tuple(coordinate.astype(int)), cv2.FONT_HERSHEY_DUPLEX, 0.3, (0, 0, 255), 1)

    cv2.imwrite(output_img_path, output_img)


def main():
    api_landmarks = pickle.load(open("/media/sdb/xzw/Adv-Makeup/Datasets_Makeup/landmark_aligned_600.pk", 'rb'))
    xzw_lmk = api_landmarks["before_aligned_600/xzw_1440.jpg"]
    id00300_lmk = api_landmarks["before_aligned_600/00300.jpg"]
    before_aligned_600_dir_path = "/media/sdb/xzw/Adv-Makeup/Datasets_Makeup/before_aligned_600/"
    landmark_align_dir_path = "/media/sdb/xzw/Adv-Makeup/Datasets_Makeup/landmark_align/"
    find_the_overlapping_idx(xzw_lmk, os.path.join(before_aligned_600_dir_path, "xzw_1440.jpg"),
                             os.path.join(landmark_align_dir_path, "xzw_added_landmark.png"))
    find_the_overlapping_idx(id00300_lmk, os.path.join(before_aligned_600_dir_path, "00300.jpg"),
                             os.path.join(landmark_align_dir_path, "00300_added_landmark.png"))


if __name__ == '__main__':
    main()
